
"""
Cognitive Map Reasoning System - Enhanced with Consolidated Data Validation
Enhanced version with unified validation functions and improved error handling
"""

import numpy as np
from typing import Union, List, Any
from numpy import ndarray, dtype


class Utils:
    @staticmethod
    def binary(x: np.ndarray) -> np.ndarray:
        """
            f(x) = 1 or 0
        """
        return (x > 0).astype(float)

    @staticmethod
    def trivalent(x: np.ndarray) -> np.ndarray:
        """
            f(x) = 1 or 0 or -1
        """
        return np.where(
            x > 0, np.array(1.0),
            np.where(
                x < 0, np.array(-1.0), np.array(0.0))
        )

    @staticmethod
    def sigmoid(x: np.ndarray, _c: float = 1.0) -> np.ndarray:
        """sigmoid transformation function with configurable slope"""
        return 1 / (1 + np.exp(-_c * x))

    @staticmethod
    def tanh(x: np.ndarray, _c: float = 1.0) -> np.ndarray:
        return np.tanh(_c * x)

    funcs_registry = {
        'binary': binary,
        'trivalent': trivalent,
        'sigmoid': sigmoid,
        'tanh': tanh,
    }

    @staticmethod
    def print_trajectory(
            _i: int,
            current_state: np.ndarray,
            weighted_state: np.ndarray,
            activated_state: np.ndarray
    ):
        print("Current iteration: ", _i)
        print("current state: ", current_state)
        print("weighted state: ", weighted_state)
        print("After activate: ", activated_state)


class CognitiveMap:
    """
    NumPy Implementation of Cognitive Map Reasoning System
    Core iteration Formula (6): x_i(t+1) = f(∑_{j≠i} w_ij * x_j(t))
    Enhanced with consolidated data validation functions
    """

    def __init__(
            self,
            num_nodes: int,
            node_names: List[str],
            init_state: np.ndarray,
            init_weight: Union[np.ndarray, np.ndarray],
            act_fn_name: str,
            max_iterations: int,
            convergence_threshold: float = 1e-6,
    ):
        # load init data
        self.num_nodes = num_nodes
        self.node_names = node_names
        self.init_state = init_state
        self.init_weight = init_weight
        self.utils = Utils()
        self.activate_function = self.utils.funcs_registry[act_fn_name]

        # control parameters
        self.max_iterations = max_iterations
        self.convergence_threshold = convergence_threshold


    def model1_iteration(self):
        # Init input parameters
        current_state = self.init_state
        current_weight = self.init_weight
        state_history = [current_state.copy()]
        diff_history = []

        for i in range(self.max_iterations):
            # 计算 权重和, 然后用tanh激活函数
            # Calculate the sum of weight and then, active with tanh
            state_update = np.dot(current_state, current_weight)
            new_state = self.activate_function(state_update)

            state_diff = np.linalg.norm(new_state - current_state)
            diff_history.append(np.max(state_diff))
            state_history.append(new_state.copy())
            current_state =  new_state    # update input state

        return np.array(state_history), np.array(diff_history)


    def model3_iteration(
            self,
            _gamma,
            _lambda
    ) -> tuple[ndarray[Any, dtype[Any]], ndarray[Any, dtype[Any]], ndarray[Any, dtype[Any]]]:
        """
            Perform cognitive map reasoning process with consolidated validation
            F1: update state
                x_i(t+1) = tanh( sum_{j=1}^{n} w_{ij}(t) * x_j(t) + γ * x_i(t) )
            F2: update weight
                w_{ij}(t+1) = tanh( λ * w_{ij}(t) + (1-λ) * (x_j(t+1) - x_i(t+1)) )
        """

        # Init input parameters
        current_state = self.init_state
        current_weight = self.init_weight

        # Append history file
        state_history = [current_state.copy()]
        weight_history = [current_weight.copy()]
        diff_history = []

        for i in range(self.max_iterations):
            # 计算 权重和 + 自反馈项, 然后用tanh激活函数
            state_update = np.dot(current_state, current_weight) + _gamma * current_state
            new_state = self.activate_function(state_update)  #

            print("input_state", "\n", current_state)
            print("state_update", "\n", state_update)
            print("new_state", "\n", new_state)
            # 计算状态差异矩阵
            state_diff_matrix = np.subtract.outer(new_state, new_state)  # x_j - x_i

            # 根据更新之后的状态 更新权重矩阵
            weight_update = _lambda * current_weight + (1 - _lambda) * state_diff_matrix
            new_weight = self.activate_function(weight_update)

            print("input_weight", "\n", current_weight)
            print("diff_matrix", "\n", state_diff_matrix)
            print("weight_update", "\n", weight_update)
            print("new_state", "\n", new_state)

            # save output state
            state_diff = np.linalg.norm(new_state - current_state)
            diff_history.append(np.max(state_diff))

            state_history.append(new_state.copy())
            weight_history.append(new_weight.copy())

            current_state =  new_state    # update input state
            current_weight = new_weight

        return np.array(state_history), np.array(weight_history), np.array(diff_history)


if __name__ == "__main__":
    n_nodes = 5
    init_vector = np.array([1.0, 0, 0, 0, 0])
    init_weights = np.array(
        [
            [0, 0, -1, 0, 1],
            [0, 0, 0, -1, 0],
            [0, -1, 0, 0, -1],
            [-1, 1, 0, 0, 0],
            [0, 0, 0, 1, 1],
        ], dtype=np.float32)

    test1 = CognitiveMap(
        num_nodes=n_nodes,
        node_names=['1', '2', '3', '4', '5'],
        init_state=init_vector,
        init_weight=init_weights,
        act_fn_name='tanh',
        max_iterations= 10
    )

    a, b, c = test1.model3_iteration(0.5, 0.5)
    print(a, b, c)

    #r_history, d_history = test1.iteration()
    #print(r_history)
    #print(d_history)
    # l=0.5
    # g=0.5
    # v, w = test1.reasoning_model2(init_vector, init_weights, g, l)


# 1.20389753 1.39802546 1.15321794 0.74420728 0.3612557
# 0.17171826 0.08829389 0.04525667 0.02286045 0.01147083





